(*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the "hack" directory of this source tree.
 *
 *)

open Hh_prelude

exception WrongEnum

module Category = struct
  [@@@warning "-32"]

  type t =
    | Decling
    | Disk_cat
    | Get_ast
    | Get_decl
    | Typecheck
  [@@deriving ord, enum]

  [@@@warning "+32"]

  let to_string (category : t) : string =
    match category with
    | Decling -> "decling"
    | Disk_cat -> "disk_cat"
    | Get_ast -> "get_ast"
    | Get_decl -> "get_decl"
    | Typecheck -> "typecheck"

  let count = (* max here is generated by ppx enum *) max + 1

  let of_enum_exn (i : int) : t =
    match of_enum i with
    | Some cat -> cat
    | None -> raise WrongEnum
end

module CategorySet = Caml.Set.Make (Category)

type time_in_sec = float

type counter = {
  count: int;  (** how many times did 'count' get called? *)
  time: time_in_sec;  (** cumulative duration of all calls to 'count' *)
}

let empty = { count = 0; time = 0. }

(** here we store each individual counter. *)
type t =
  (* Making this a map causes 1% typing time regression. *)
  counter Array.t

let counters : t ref = ref (Array.create ~len:Category.count empty)

let restore_state (new_state : t) : unit = counters := new_state

let reset () : t =
  let old_counters = !counters in
  counters := Array.create ~len:Category.count empty;
  old_counters

let get_counter (category : Category.t) : counter =
  !counters.(Category.to_enum category)

let set_counter (category : Category.t) (counts : counter) : unit =
  !counters.(Category.to_enum category) <- counts

let get_time (category : Category.t) : unit -> time_in_sec =
  match category with
  | Category.Disk_cat ->
    (* Wall-clock time, implemented in C:
       let t = gettimeofday() in t.tv_sec + t.tv_usec / 1e6 *)
    Unix.gettimeofday
  | Category.Typecheck
  | Category.Decling
  | Category.Get_ast
  | Category.Get_decl ->
    (* CPU time, excluding I/O, implemented in C in one of three ways depending on ocaml compilation flags:
       1. let r = getrusage(RUSAGE_SELF) in r.ru_utime.tv_sec + r.ru_utime.tv_usec / 1e6
       2. let t = times() in t.tms_utime + t.tms_stime
       3. clock() / CLOCKS_PER_SPEC *)
    Sys.time

let get_time_label (category : Category.t) : string =
  match category with
  | Category.Disk_cat -> "walltime"
  | Category.Typecheck
  | Category.Decling
  | Category.Get_ast
  | Category.Get_decl ->
    "cputime"

let count (category : Category.t) (f : unit -> 'a) : 'a =
  let get_time = get_time category in
  let tally = get_counter category in
  let start_time = get_time () in
  Utils.try_finally ~f ~finally:(fun () ->
      set_counter
        category
        {
          count = tally.count + 1;
          time = tally.time +. get_time () -. start_time;
        })

let read_time (category : Category.t) : time_in_sec =
  (get_counter category).time

let get_counters () : Telemetry.t =
  let telemetry_of_counter category counter =
    Telemetry.create ()
    |> Telemetry.int_ ~key:"count" ~value:counter.count
    |> Telemetry.float_ ~key:(get_time_label category) ~value:counter.time
  in
  let telemetry =
    Array.foldi
      ~f:(fun i telemetry counter ->
        let category = Category.of_enum_exn i in
        let telemetry =
          Telemetry.object_
            telemetry
            ~key:(Category.to_string category)
            ~value:(telemetry_of_counter category counter)
        in
        telemetry)
      !counters
      ~init:(Telemetry.create ())
  in
  telemetry
